#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
模型管理器
负责模型和必要数据的下载、管理和加载
"""

import os
import sys
import logging
import shutil
from pathlib import Path

# 添加项目根目录到路径
sys.path.append(str(Path(__file__).resolve().parent.parent))

# 导入项目模块
from config.paths import (
    MODELS_DIR, TORCH_MODELS_DIR, DATA_DIR, NLTK_DATA_DIR,
    GDRIVE_FOLDER_ID
)

# 配置日志
logger = logging.getLogger("模型管理器")

class ModelManager:
    """模型和数据管理器类"""
    
    def __init__(self, ui_callback=None):
        """
        初始化模型管理器
        
        Args:
            ui_callback: UI回调函数，用于在界面显示下载状态
        """
        self.ui_callback = ui_callback
    
    def download_dependencies(self, force=False):
        """
        下载必要的依赖项和模型
        
        Args:
            force: 是否强制重新下载
            
        Returns:
            bool: 是否成功
        """
        success = True
        
        # 第1步：下载NLTK数据
        if self._download_nltk_data(force) is False:
            success = False
            
        # 第2步：下载模型文件
        if self._download_models(force) is False:
            success = False
            
        return success
    
    def _download_nltk_data(self, force=False):
        """
        下载NLTK数据
        
        Args:
            force: 是否强制重新下载
            
        Returns:
            bool: 是否成功
        """
        try:
            # 检查NLTK数据是否已存在
            punkt_dir = NLTK_DATA_DIR / "tokenizers" / "punkt"
            if punkt_dir.exists() and not force:
                logger.info("NLTK punkt数据已存在，跳过下载")
                return True
                
            # 导入NLTK
            import nltk
            
            # 设置NLTK数据目录
            nltk.data.path = [str(NLTK_DATA_DIR)]
            
            # 进度显示
            if self.ui_callback:
                self.ui_callback("正在下载NLTK punkt数据...")
                
            # 下载数据
            logger.info("开始下载NLTK punkt数据...")
            nltk.download('punkt', download_dir=str(NLTK_DATA_DIR))
            
            # 检查是否成功
            if punkt_dir.exists():
                logger.info("NLTK punkt数据下载成功")
                return True
            else:
                logger.error("NLTK punkt数据下载失败")
                return False
        except Exception as e:
            logger.error(f"下载NLTK数据时出错: {str(e)}")
            return False
    
    def _download_models(self, force=False):
        """
        下载模型文件
        
        Args:
            force: 是否强制重新下载
            
        Returns:
            bool: 是否成功
        """
        try:
            # 检查模型文件是否已存在
            acoustic_model = TORCH_MODELS_DIR / "acoustic_model.pt"
            style_predictor = TORCH_MODELS_DIR / "style_predictor.pt"
            vocoder = TORCH_MODELS_DIR / "vocoder.pt"
            
            if (acoustic_model.exists() and 
                style_predictor.exists() and 
                vocoder.exists() and 
                not force):
                logger.info("模型文件已存在，跳过下载")
                return True
                
            # 导入gdown
            try:
                import gdown
            except ImportError:
                logger.error("未安装gdown，无法下载模型文件")
                return False
                
            # 创建模型目录
            os.makedirs(TORCH_MODELS_DIR, exist_ok=True)
            
            # 进度显示
            if self.ui_callback:
                self.ui_callback("正在下载模型文件...")
                
            # 下载文件夹
            logger.info(f"开始从Google Drive下载模型文件 (folder_id: {GDRIVE_FOLDER_ID})...")
            try:
                gdown.download_folder(
                    f"https://drive.google.com/drive/folders/{GDRIVE_FOLDER_ID}",
                    output=str(TORCH_MODELS_DIR),
                    use_cookies=False
                )
            except Exception as e:
                logger.error(f"使用gdown下载模型时出错: {str(e)}")
                return False
                
            # 检查是否成功
            if (acoustic_model.exists() and 
                style_predictor.exists() and 
                vocoder.exists()):
                logger.info("模型文件下载成功")
                return True
            else:
                logger.error("模型文件下载不完整")
                return False
        except Exception as e:
            logger.error(f"下载模型文件时出错: {str(e)}")
            return False
    
    def check_model_files(self):
        """
        检查模型文件是否存在
        
        Returns:
            bool: 是否所有必要的模型文件都存在
        """
        required_files = [
            TORCH_MODELS_DIR / "acoustic_model.pt",
            TORCH_MODELS_DIR / "style_predictor.pt",
            TORCH_MODELS_DIR / "vocoder.pt"
        ]
        
        for file_path in required_files:
            if not file_path.exists():
                logger.warning(f"缺少模型文件: {file_path}")
                return False
        
        logger.info("所有必要的模型文件都存在")
        return True
    
    def cleanup_temp_files(self):
        """清理临时文件"""
        pass

# 示例用法
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    manager = ModelManager()
    manager.download_dependencies() 